# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: chunhui

import json
import os
import pandas as pd
import shutil
import subprocess
import tempfile
import time
from absl import app, flags
from coin.exchange.ftx_futures.kr_rest.native_private_client import \
   FTXFuturesNativePrivateClient


FLAGS = flags.FLAGS


class FTXSwapClient(FTXFuturesNativePrivateClient):
  def __init__(self, key_file):
    super().__init__(key_file=key_file)
    with open('data/coin2/product_info/Futures.Ftx.v1.json') as f:
      product_info = json.load(f)
    self.mea = product_info['mea']
    self.product_info = product_info['product_infos']

  def convert_native_symbol(self, symbol):
    for info in self.product_info:
      if info['symbol'] == symbol:
        return info['native_symbol']
    raise ValueError(f'Unrecognized symbol {symbol}')

  def _set_leverage(self, symbol, leverage):
    method = 'POST'
    path = '/api/v5/account/set-leverage'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        'instId': native_symbol,
        'lever': leverage,
        'mgnMode': 'cross',
        'posSide': 'net'
    }
    update = self._query(method, path, params)
    return update

  def set_leverage(self, symbols, leverage):
    if isinstance(symbols, str):
      symbols = symbols.split(',')
    for i, symbol in enumerate(symbols):
      if (i + 1) % 5 == 0:
        time.sleep(2)
      result = self._set_leverage(symbol, leverage)
      print(result)

  def _get_pos_limit(self, symbol):
    method = 'GET'
    path = '/api/v5/account/max-size'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "tdMode": "cross",
        "ccy": "USDT"
    }
    update = self._query(method, path, params)
    return update

  def _get_trade_limit(self, symbol):
    method = 'GET'
    path = '/api/v5/account/max-avail-size'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "tdMode": "cross",
        "ccy": "USDT"
    }
    update = self._query(method, path, params)
    return update

  def _get_leverage(self, symbol):
    method = 'GET'
    path = '/api/v5/account/leverage-info/'
    native_symbol = self.convert_native_symbol(symbol)
    params = {
        "instId": native_symbol,
        "mgnMode": "cross"
    }
    update = self._query(method, path, params)
    return update

  def get_leverage(self, symbols):
    if isinstance(symbols, str):
      symbols = symbols.split(',')
    data = []
    for i, symbol in enumerate(symbols):
      if (i + 1) % 19 == 0:
        time.sleep(2)
      leverage = self._get_leverage(symbol)
      pos_limit = self._get_pos_limit(symbol)
      trade_limit = self._get_trade_limit(symbol)
      data.extend(leverage.as_json()['msg']['data'])
      data[-1].update(pos_limit.as_json()['msg']['data'][0])
      data[-1].update(trade_limit.as_json()['msg']['data'][0])
    df = pd.DataFrame(data)
    return df

  def _test(self):
    method = 'GET'
    path = '/account'
    import pdb; pdb.set_trace()
    update = self._query(method, path)
    import pdb; pdb.set_trace()
    return update

  def test(self):
    self._test()


def init_flags():
  flags.DEFINE_string('strategy', '', 'strategy_name')
  flags.DEFINE_string('account', '', 'account name')
  flags.DEFINE_string('symbol', 'BTC-USDT.PERPETUAL', 'usdt amount')
  flags.DEFINE_integer('leverage', 10, 'upd leverage')
  flags.DEFINE_bool('query_only', True, 'check leverage')


def main(_):
  pd.set_option('display.max_rows', None)
  account = flags.FLAGS.account
  strat_symbols = [
    "1INCH-USD.PERPETUAL",
    "AAVE-USD.PERPETUAL",
    "ALICE-USD.PERPETUAL",
    "ALPHA-USD.PERPETUAL",
    "AUDIO-USD.PERPETUAL",
    "AXS-USD.PERPETUAL",
    "BADGER-USD.PERPETUAL",
    "BCHN-USD.PERPETUAL",
    "BNB-USD.PERPETUAL",
    "BTC-USD.PERPETUAL",
    "C98-USD.PERPETUAL",
    "CHZ-USD.PERPETUAL",
    "COMP-USD.PERPETUAL",
    "CRV-USD.PERPETUAL",
    "DOGE-USD.PERPETUAL",
    "ENJ-USD.PERPETUAL",
    "ETH-USD.PERPETUAL",
    "FTM-USD.PERPETUAL",
    "FTT-USD.PERPETUAL",
    "GRT-USD.PERPETUAL",
    "HNT-USD.PERPETUAL",
    "KIN-USD.PERPETUAL",
    "LINK-USD.PERPETUAL",
    "LTC-USD.PERPETUAL",
    "MATIC-USD.PERPETUAL",
    "MKR-USD.PERPETUAL",
    "MTL-USD.PERPETUAL",
    "OMG-USD.PERPETUAL",
    "PERP-USD.PERPETUAL",
    "RAY-USD.PERPETUAL",
    "REN-USD.PERPETUAL",
    "RSR-USD.PERPETUAL",
    "RUNE-USD.PERPETUAL",
    "SKL-USD.PERPETUAL",
    "SLP-USD.PERPETUAL",
    "SNX-USD.PERPETUAL",
    "SOL-USD.PERPETUAL",
    "SRM-USD.PERPETUAL",
    "STEP-USD.PERPETUAL",
    "SUSHI-USD.PERPETUAL",
    "SXP-USD.PERPETUAL",
    "TOMO-USD.PERPETUAL",
    "TRU-USD.PERPETUAL",
    "TRX-USD.PERPETUAL",
    "UNI-USD.PERPETUAL",
    "USDT-USD.PERPETUAL",
    "WAVES-USD.PERPETUAL",
    "XRP-USD.PERPETUAL",
    "YFI-USD.PERPETUAL",
    "YFII-USD.PERPETUAL"
  ]

  tmp_dir = tempfile.mkdtemp()
  try:
    key_file = os.path.join(tmp_dir, 'key.json')
    cmd = f"""./python/coin_deploy/fab.sh \
-f python/coin_deploy/fabhelper/fabutil.py \
-P upload_coin_key_high:vault_token_file="./python/coin_deploy/vault/coin_key_high.json",\
key_name="trade/{account}/FTX/trade_key.json",\
remote_key_path="{key_file}",\
run_local=True
"""
    _ = subprocess.check_output(cmd, shell=True)
    msc = FTXSwapClient(key_file=key_file)
    msc.test()

    result = msc.get_leverage(strat_symbols)
    print("\n#####################")
    print("-- Before Adjust --")
    print(result)
    print("######################\n")

    if not flags.FLAGS.query_only:
      assert flags.FLAGS.symbol
      msc.set_leverage(FLAGS.symbol, FLAGS.leverage)
      print("\n#####################")
      print("-- After Adjust --")
      print(msc.get_leverage(strat_symbols))
      print("######################\n")
  except Exception:
    raise
  finally:
    shutil.rmtree(tmp_dir)


'''
./pyrunner cc/appcoin2/strategy/basis_strat/support/adjust_leverage/ftx_leverage.py
--strategy=mmrev_okex_perpetual_1
--query_only=0
--leverage=10
--symbol=BTC-USDT.PERPETUAL,ETH-USDT.PERPETUAL
'''

if __name__ == '__main__':
  init_flags()
  app.run(main)
